import torch
import torchvision
import numpy as np
import sys
sys.path.append("..") # 为了导⼊上层⽬录的d2lzh_pytorch
import d2lzh_pytorch as d2l
import matplotlib
matplotlib.use('TkAgg')  # 或者 'Qt5Agg'
import matplotlib.pyplot as plt


#设置批量
batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

#28*28像素点得到784个数据输入
num_inputs = 784
num_outputs = 10
W = torch.tensor(np.random.normal(0, 0.01, (num_inputs, num_outputs)), dtype=torch.float)
b = torch.zeros(num_outputs, dtype=torch.float)

#求模型的参数梯度
W.requires_grad_(requires_grad=True)
b.requires_grad_(requires_grad=True)


def softmax(X):
    X_exp = X.exp()    #这里使得每个数的特征更加明显，同时保证为非负数
    partition = X_exp.sum(dim=1, keepdim=True)
    return X_exp / partition # 这⾥应⽤了⼴播机制

#定义softmax回归模型
def net(X):
    return softmax(torch.mm(X.view((-1, num_inputs)), W) + b)

#定义交叉熵损失函数
def cross_entropy(y_hat, y):
    return - torch.log(y_hat.gather(1, y.view(-1, 1)))

#计算分类的准确率
def accuracy(y_hat, y):
    return (y_hat.argmax(dim=1) == y).float().mean().item()

def show_fashion_mnist(images, labels):
    # 使用合适的显示后端
    use_interactive_display()  # 如果需要矢量图，保持此行，或改为交互式后端显示

    # 创建子图并设置图形大小
    _, figs = plt.subplots(1, len(images), figsize=(12, 12))

    # 遍历每个子图，显示图像及标签
    for f, img, lbl in zip(figs, images, labels):
        # 显示28x28的图片，并设置标题为标签
        f.imshow(img.view((28, 28)).numpy(), cmap='gray')  # 显示为灰度图
        f.set_title(lbl)  # 设置标题
        f.axes.get_xaxis().set_visible(False)  # 隐藏x轴
        f.axes.get_yaxis().set_visible(False)  # 隐藏y轴

    plt.show()  # 显示图形

def use_interactive_display():
    # 使用交互式后端显示
    matplotlib.use('TkAgg')  # 或者使用 'Qt5Agg'，视你的环境而定

num_epochs, lr = 5, 0.1
d2l.train_softmax(net, train_iter, test_iter, cross_entropy, num_epochs, batch_size, [W, b], lr)

# 获取一个批次数据
X, y = next(iter(test_iter))
# 获取真实标签和预测标签
true_labels = d2l.get_fashion_mnist_labels(y.numpy())  # 获取真实标签
pred_labels = d2l.get_fashion_mnist_labels(net(X).argmax(dim=1).numpy())  # 获取预测标签
# 拼接真实标签和预测标签
titles = [true + '\n' + pred for true, pred in zip(true_labels, pred_labels)]
# 显示前9张图片和标签
show_fashion_mnist(X[0:9], titles[0:9])
